
from ComputerVision.searchLaserPoint import searchRedLaserPoint,METHOD_HSV_RANGE2,METHOD_HSV_RANGE3
from Drivers.FOCMotorController import FOCMotorController
from Algorithm.PID import PID
from picamera2 import Picamera2

import cv2 as cv
import numpy as np
import time
import os
import threading
from RPi import GPIO

os.environ['QT_QPA_PLATFORM'] = 'wayland'

GPIO.setmode(GPIO.BCM)
GPIO.setup(17,GPIO.OUT)
GPIO.output(17,GPIO.HIGH)

picam = Picamera2()
picamConfig = picam.create_video_configuration(main={"size": (1920,1080),'format':'BGR888'}, lores={"size": (640, 480)}, display="main")
picam.configure(picamConfig)
picam.set_controls({"ExposureTime":20000})
picam.start()

globalLock = threading.Lock()
TASK_RUNNING = False

def drawPoint(
    img: np.ndarray,
    point: list,
    color: list = [255,0,0],
    radius: int = 5
    ):
    img = cv.circle(img,np.int32(point),radius,color,-1)
    return img

def getImage():
    return cv.cvtColor(
        picam.capture_array(),
        cv.COLOR_RGB2BGR
    )
def cropToRequiredSize(img):
    return img[:,:]

yawMotor = FOCMotorController('/dev/serial/by-id/usb-STMicroelectronics_QDrive_ComPort_2068358E4831-if00')
pinchMotor = FOCMotorController('/dev/serial/by-id/usb-STMicroelectronics_QDrive_ComPort_205F3298484E-if00')

def relativeAngle(p1,p0):
    # 计算相对角度
    return math.atan2(p1[1]-p0[1],p1[0]-p0[0])
def relativeDistance(p1,p0):
    # 计算相对距离
    return ((p1[1]-p0[1])**2 + (p1[0]-p0[0])**2)**0.5

def setTaskRunningState(state):
    global TASK_RUNNING
    globalLock.acquire()
    TASK_RUNNING = state
    globalLock.release()

def task1():
    global TASK_RUNNING
    yawPID = PID(0.012,0.0,0.0,1)
    pinchPID = PID(-0.012,0.0,0.0,1)
    # 绿色激光点坐标，一开始写死
    greenLaserPoint = [940,475] 
    # 单次追踪激光点
    while True:
        if TASK_RUNNING:
            image = getImage()
            image = cropToRequiredSize(image)
            redLaserPoint = searchRedLaserPoint(image,method=METHOD_HSV_RANGE2)
            if not redLaserPoint:
                redLaserPoint = searchRedLaserPoint(image,method=METHOD_HSV_RANGE3)
            if redLaserPoint:
                xLoss = greenLaserPoint[0] - redLaserPoint[0]
                yLoss = greenLaserPoint[1] - redLaserPoint[1]
                yawFeedback = yawPID.forward(xLoss,0)
                pinchFeedback = pinchPID.forward(yLoss,0)
                globalLock.acquire()
                yawMotor.setSpeed(yawFeedback,lowSpeedMode=True)
                pinchMotor.setSpeed(pinchFeedback,lowSpeedMode=True)
                globalLock.release()
                print(xLoss,yLoss)
                if relativeDistance([xLoss,yLoss],[0,0]) < 10:
                    # time.sleep(0.5)
                    globalLock.acquire()
                    yawMotor.setSpeed(0,lowSpeedMode=True)
                    pinchMotor.setSpeed(0,lowSpeedMode=True)
                    globalLock.release()
                    GPIO.output(17,GPIO.LOW)
                    time.sleep(1)
                    GPIO.output(17,GPIO.HIGH)
                    break
        else:
            return
    setTaskRunningState(False)
            
            
    
def task2():
    global TASK_RUNNING
    yawPID = PID(
        0.012,
        0.0,
        0.0,
        1
        )
    pinchPID = PID(
        -0.012,
        0.0,
        0.0,
        1
        )
    # 绿色激光点坐标，一开始写死
    greenLaserPoint = [940,475] 
    # 连续追踪激光点
    while True:
        if TASK_RUNNING:
            image = getImage()
            image = cropToRequiredSize(image)
            redLaserPoint = searchRedLaserPoint(image,method=METHOD_HSV_RANGE2)
            if not redLaserPoint:
                redLaserPoint = searchRedLaserPoint(image,method=METHOD_HSV_RANGE3)
            if redLaserPoint:
                xLoss = greenLaserPoint[0] - redLaserPoint[0]
                yLoss = greenLaserPoint[1] - redLaserPoint[1]
                yawFeedback = yawPID.forward(xLoss,0)
                pinchFeedback = pinchPID.forward(yLoss,0)
                globalLock.acquire()
                yawMotor.setSpeed(yawFeedback,lowSpeedMode=True)
                pinchMotor.setSpeed(pinchFeedback,lowSpeedMode=True)
                globalLock.release()
                print(xLoss,yLoss)
                if relativeDistance([xLoss,yLoss],[0,0]) < 75:
                    globalLock.acquire()
                    GPIO.output(17,GPIO.LOW)
                    globalLock.release()
                else:
                    globalLock.acquire()
                    GPIO.output(17,GPIO.HIGH)
                    globalLock.release()
            else:
                globalLock.acquire()
                GPIO.output(17,GPIO.HIGH)
                globalLock.release()
        else:
            globalLock.acquire()
            yawMotor.setSpeed(0,lowSpeedMode=True)
            pinchMotor.setSpeed(0,lowSpeedMode=True)
            globalLock.release()
            return
        # if redLaserPoint:
        #     image = drawPoint(image,redLaserPoint)
        # cv.imshow('task2',image)
        # k = cv.waitKey(1)
        # if k == ord('q'):
        #     yawMotor.setSpeed(0)
        #     pinchMotor.setSpeed(0)
        #     break
    # cv.destroyAllWindows()

def runTask1():
    global TASK_RUNNING
    if TASK_RUNNING == True:
        print('有任务正在运行')
        return
    else:
        setTaskRunningState(True)
        t = threading.Thread(target=task1)
        t.start()

def runTask2():
    global TASK_RUNNING
    if TASK_RUNNING == True:
        print('有任务正在运行')
        return
    else:
        setTaskRunningState(True)
        t = threading.Thread(target=task2)
        t.start()

def taskEmergeStop():
    # 急停
    global TASK_RUNNING
    globalLock.acquire()
    TASK_RUNNING = False
    GPIO.output(17,GPIO.HIGH)
    yawMotor.setSpeed(0,lowSpeedMode=True)
    pinchMotor.setSpeed(0,lowSpeedMode=True)
    globalLock.release()

def initMotor():
    yawMotor.enable()
    pinchMotor.enable()
    yawMotor.setSpeed(0)
    pinchMotor.setSpeed(0)

if __name__ == '__main__':
    yawMotor.enable()
    pinchMotor.enable()
    yawMotor.setSpeed(0)
    pinchMotor.setSpeed(0)
    print('===============')
    print('2023电赛E题2号机程序')
    try:
        while True:
            print('(1)任务:追踪激光点（单次）\n(2)任务:连续追踪激光点\n(6)使能电机\n(7)禁用电机\n(0)退出')
            code = input('请输入任务编号>')
            
            if code == '1':
                task1()
            elif code == '2':
                task2()
            elif code == '6':
                yawMotor.enable()
                pinchMotor.enable()
            elif code == '7':
                yawMotor.disable()
                pinchMotor.disable()
            elif code == '0':
                break
    except KeyboardInterrupt:
        yawMotor.disable()
        pinchMotor.disable()
        GPIO.output(17,GPIO.HIGH)
        pass
    